from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVC
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from Tools import readbunchobj, metrics_result
import numpy as np
from joblib import dump

# 导入训练集和测试集
trainpath = "../train_word_bag/tf_idf_space.dat"
train_set = readbunchobj(trainpath)
testpath = "../test_word_bag/test_tf_idf_space.dat"
test_set = readbunchobj(testpath)

# 输出单词矩阵的类型
print("标准化/归一化前的矩阵形状:")
print(np.shape(train_set.tdm))
print(np.shape(test_set.tdm))

# 创建一个管道，包括标准化和SVM分类器
pipeline = Pipeline([
    ('scaler', StandardScaler(with_mean=False)),
    ('svm', SVC(kernel='linear', random_state=0))
])

# 定义要搜索的参数网格
param_grid = {
    'svm__C': [0.1, 1, 10, 100],
    # 添加其他你想尝试的超参数
}

# 使用GridSearchCV进行并行搜索
grid_search = GridSearchCV(pipeline, param_grid, cv=5, n_jobs=-1, verbose=1)

# 训练模型
print("开始训练支持向量机模型...")
grid_search.fit(train_set.tdm, train_set.label)
print("支持向量机模型训练完成。")

# 保存最佳模型
best_model = grid_search.best_estimator_
dump(best_model, 'best_SVM_model.joblib')

# 预测分类结果
print("USE: 支持向量机模型 预测分类结果")
SVM_predicted = best_model.predict(test_set.tdm)
SVM_total = len(SVM_predicted)
print("结束")

# 性能评估
print("\n支持向量机模型：")
metrics_result(test_set.label, SVM_predicted, SVM_total)
